# -*- utf-8 -*-

import jieba
import os
import json

jieba.load_userdict("./distDict.txt")
def loadfile():
    combined = []

    train_text = "./fasttext_train.text"
    test_text = "./fasttext_test.text"
    remove_workd = [',', '：', "、", "，", "的", "\r", "\n"]
    for key in range(1, 10):
        remove_workd.append(str(key)+".")

    for parent, dirnames, filenames in os.walk("./data"):
        for filename in filenames:
            filename = os.path.join(parent, filename)
            with open(filename, "r+") as f:
                qd_infos = json.loads(f.read())["qd_infos"]
                for index, qd in enumerate(qd_infos):
                    texts = [qd['name'], qd['spec'], qd['raw_paragraph']]
                    childs = qd['children']
                    for child in childs:
                        texts.append(child['name'])
                        materials = child['material']
                        for material in materials:
                            texts.append(material['name'])
                    category = qd['category']
                    category = "/".join(category.split("/")[1:])
                    texts = " ".join(texts)
                    for word in remove_workd:
                        texts = texts.replace(word, " ")
                    combined.append(" ".join(jieba.lcut(texts)) + " __label__"+category + "\n")

    train_data = []
    test_data = []
    for index, line in enumerate(combined):
        if float(index) / len(combined) > 0.8:
            test_data.append(line)
        else:
            train_data.append(line)

    with open(train_text, "w") as f:
        f.writelines(train_data)

    with open(test_text, "w") as f:
        f.writelines(test_data)


if __name__ == "__main__":
    loadfile()